"""
K线图表模块 - 提供K线图生成功能
"""

import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
import numpy as np
from typing import List, Dict, Any, Optional

def create_candlestick_chart(
    df: pd.DataFrame, 
    title: str = "股票价格走势", 
    height: int = 600,
    show_volume: bool = True,
    ma_periods: List[int] = [5, 10, 20]
) -> go.Figure:
    """
    创建K线图
    
    参数:
        df: 包含OHLC数据的DataFrame，必须包含 '日期', '开盘', '最高', '最低', '收盘' 列
        title: 图表标题
        height: 图表高度
        show_volume: 是否显示成交量
        ma_periods: 移动平均周期列表
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    required_cols = ['日期', '开盘', '最高', '最低', '收盘']
    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要列: {col}")
    
    # 计算技术指标
    for period in ma_periods:
        df[f'MA{period}'] = df['收盘'].rolling(window=period).mean()
    
    # 创建K线图
    fig = go.Figure()
    
    # 添加K线
    fig.add_trace(go.Candlestick(
        x=df['日期'],
        open=df['开盘'],
        high=df['最高'],
        low=df['最低'],
        close=df['收盘'],
        name='K线',
        increasing_line_color='red',
        decreasing_line_color='green'
    ))
    
    # 添加均线
    colors = ['blue', 'orange', 'purple', 'cyan', 'magenta']
    for i, period in enumerate(ma_periods):
        color = colors[i % len(colors)]
        fig.add_trace(go.Scatter(
            x=df['日期'], 
            y=df[f'MA{period}'], 
            mode='lines', 
            name=f'MA{period}', 
            line=dict(color=color)
        ))
    
    # 设置图表布局
    fig.update_layout(
        title=title,
        xaxis_title="日期",
        yaxis_title="价格",
        height=height,
        xaxis_rangeslider_visible=False,
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1
        )
    )
    
    # 如果需要显示成交量并且数据中包含成交量列
    if show_volume and '成交量' in df.columns:
        # 创建成交量图表
        volume_fig = create_volume_chart(df, title=f"{title} - 成交量")
        
        # 创建子图布局
        fig = go.Figure()
        
        # 添加K线子图
        fig.add_trace(go.Candlestick(
            x=df['日期'],
            open=df['开盘'],
            high=df['最高'],
            low=df['最低'],
            close=df['收盘'],
            name='K线',
            increasing_line_color='red',
            decreasing_line_color='green'
        ))
        
        # 添加均线
        for i, period in enumerate(ma_periods):
            color = colors[i % len(colors)]
            fig.add_trace(go.Scatter(
                x=df['日期'], 
                y=df[f'MA{period}'], 
                mode='lines', 
                name=f'MA{period}', 
                line=dict(color=color)
            ))
        
        # 确保数据类型一致
        df_copy = df.copy()
        for col in ['开盘', '收盘', '最高', '最低']:
            if col in df_copy.columns:
                df_copy[col] = pd.to_numeric(df_copy[col], errors='coerce')
        
        # 添加成交量子图
        try:
            fig.add_trace(go.Bar(
                x=df_copy['日期'],
                y=df_copy['成交量'],
                name='成交量',
                marker=dict(
                    color=np.where(df_copy['收盘'] >= df_copy['开盘'], 'red', 'green')
                ),
                yaxis="y2"
            ))
        except Exception as e:
            # 如果比较失败，使用默认颜色
            print(f"警告: 无法比较价格数据: {str(e)}")
            fig.add_trace(go.Bar(
                x=df['日期'],
                y=df['成交量'],
                name='成交量',
                yaxis="y2"
            ))
        
        # 设置子图布局
        fig.update_layout(
            title=title,
            xaxis_title="日期",
            yaxis_title="价格",
            height=height,
            xaxis_rangeslider_visible=False,
            legend=dict(
                orientation="h",
                yanchor="bottom",
                y=1.02,
                xanchor="right",
                x=1
            ),
            yaxis=dict(
                domain=[0.3, 1]
            ),
            yaxis2=dict(
                domain=[0, 0.2],
                title="成交量"
            )
        )
    
    return fig

def create_volume_chart(df: pd.DataFrame, title: str = "成交量", height: int = 300) -> go.Figure:
    """
    创建成交量图表
    
    参数:
        df: 包含成交量数据的DataFrame，必须包含 '日期', '成交量', '开盘', '收盘' 列
        title: 图表标题
        height: 图表高度
    
    返回:
        plotly图表对象
    """
    # 确保数据格式正确
    required_cols = ['日期', '成交量', '开盘', '收盘']
    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"数据缺少必要列: {col}")
    
    # 确保数据类型一致
    try:
        # 尝试将价格列转换为float类型
        df_copy = df.copy()
        for col in ['开盘', '收盘', '最高', '最低']:
            if col in df_copy.columns:
                df_copy[col] = pd.to_numeric(df_copy[col], errors='coerce')
        
        # 创建成交量图表
        fig = px.bar(
            df_copy,
            x='日期',
            y='成交量',
            color=np.where(df_copy['收盘'] >= df_copy['开盘'], 'up', 'down'),
            color_discrete_map={'up': 'red', 'down': 'green'},
            title=title
        )
    except Exception as e:
        # 如果转换失败，使用默认颜色
        print(f"警告: 无法比较价格数据: {str(e)}")
        fig = px.bar(
            df,
            x='日期',
            y='成交量',
            title=title
        )
    
    # 设置图表布局
    fig.update_layout(
        height=height,
        xaxis_title="日期",
        yaxis_title="成交量",
        showlegend=False
    )
    
    return fig